Summary

• Bug: _capture_generation_cuda_graphs in PyTorchModelEngine lacks the cross-rank batch-None check that _general_warmup_impl and _run_autotuner_warmup already have.
• Impact: DEP mode + enable_attention_dp=true + CUDA graph batch sizes ≥ 16 → permanent distributed deadlock during engine warmup. Process hangs before producing any Python output (run.log = 0 bytes), 100% reproducible.
• Fix: Mirror the _assert_all_tp_ranks_have_warmup_batch guard (already present in the two other warmup paths) into _capture_generation_cuda_graphs.

Root Cause

_capture_generation_cuda_graphs iterates batch sizes in reverse order (largest first, so smaller graphs can reuse the memory pool). Under attention-DP, KV-cache capacity differs per TP rank. _create_cuda_graph_warmup_request returns None on ranks that lack space. Without a cross-rank check:

• Ranks where batch is None silently continue
• Other ranks enter forward() containing tp_comm collectives (NCCL allreduce / DEP alltoall)
• Skipped ranks never reach the collective → permanent deadlock

The two other warmup paths already guard this with _assert_all_tp_ranks_have_warmup_batch (added in an earlier PR). This PR closes the remaining gap.

Fix

# Before
if batch is None:
    continue

# After
if batch is None and self.mapping.tp_size <= 1:
    continue
self._assert_all_tp_ranks_have_warmup_batch(batch, bs)
if batch is None:
    continue  # all ranks agree: skip

Tests

Added tests/unittest/_torch/executor/test_cuda_graph_capture_adp_guard.py with four unit tests:

1. Structural: guard call is present in source (regression sentinel)
2. Asymmetric None raises: some ranks None, others valid → RuntimeError before forward()
3. All None skips: all ranks None → graceful skip, no raise
4. All valid proceeds: all ranks valid → forward() called normally

Test plan

• Run pytest tests/unittest/_torch/executor/test_cuda_graph_capture_adp_guard.py (no GPU needed)
  → 4/4 passed (no GPU required; 1.61 s wall time on 8×B300 host)
• End-to-end: DEP + enable_attention_dp=true + cuda_graph_config.batch_sizes=[1,2,...,32] no longer deadlocks at warmup
  → Verified via DSv4 Pro DEP GSM8K eval (8×B300, MTP=3, GVR ON, BS up to 32):
  engine initialized cleanly in 452 s, inferred all 1319 problems without hang.
  Accuracy: 96.51% average (flexible-extract 96.51, strict-match 96.51),
  matching TEP mode (96.63%) within ±0.5 pp statistical error.

Summary by CodeRabbit

Release Notes

• Bug Fixes

  • Improved CUDA graph generation to detect and handle asymmetric KV-cache capacity across distributed ranks, preventing potential deadlocks in multi-rank configurations.

• Tests

  • Added comprehensive unit tests validating correct behavior under asymmetric batch conditions in distributed tensor-parallel execution scenarios.